from os.path import join, abspath
from typing import List

from backend.experiment.framework.shots_prompt.task import \
    PerplexityPromptShotsClassificationTaskMaker
from backend.experiment.framework.task import TaskMaker
from backend.experiment.tokenizer import tokenizer
from backend.experiment.framework.shots_prompt.implement import \
    CommonStaticShotsGenerator, \
    CommonPromptClassificationSample
from backend.experiment.model import context_limit
from backend.experiment.clue.ocnli.scheme0.implement import OCNLIImplement as S0


class OCNLIImplement(S0):
    @classmethod
    def _make_test_sample(cls, record, made_sample_num: int) -> \
            CommonPromptClassificationSample or \
            List[CommonPromptClassificationSample]:
        prompt = cls._make_prompt(record, True)
        return CommonPromptClassificationSample(
            prompt=prompt,
            target_mask_token=cls.__mask_token,
            sample_id=made_sample_num
        )

    def get_task_maker(self) -> TaskMaker:
        return PerplexityPromptShotsClassificationTaskMaker(
            token_limit=context_limit,
            tokenizer=tokenizer,
            shots_generator=CommonStaticShotsGenerator(
                example_tokens=self.get_shuffle_example_tokens()
            ),
            label_map=self.label_map()
        )

    @classmethod
    def cache_dir(cls) -> str:
        return join(cls.work_dir(), 'cache')

    @classmethod
    def work_dir(cls) -> str:
        return abspath(join(__file__, '..'))
